#!/usr/bin/env python3

import rospy
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import os
import glob
import time


class VirtualCamera:
    def __init__(self):
        # 初始化ROS节点
        rospy.init_node('virtual_camera', anonymous=False)
        
        # 创建CV Bridge用于图像转换
        self.bridge = CvBridge()
        
        # 获取ROS参数
        self.is_video_mode = rospy.get_param('~is_video_mode', False)  # 默认为图像模式
        self.data_path = rospy.get_param('~data_path', "test_data/images")     # 数据路径
        self.fps = rospy.get_param('~fps', 60)                     # 发布频率
        self.loop = rospy.get_param('~loop', False)                   # 是否循环播放
        
        # 发布图像话题
        self.image_pub = rospy.Publisher('/camera/image_raw', Image, queue_size=10)
        
        # 初始化数据
        self.current_index = 0
        self.data_list = []
        self.video_cap = None
        
        # 设置发布频率
        self.rate = rospy.Rate(self.fps)
        
        rospy.loginfo(f"Virtual Camera initialized:")
        rospy.loginfo(f"  Mode: {'Video' if self.is_video_mode else 'Image'}")
        rospy.loginfo(f"  Data path: {self.data_path}")
        rospy.loginfo(f"  FPS: {self.fps}")
        rospy.loginfo(f"  Loop: {self.loop}")
        
        # 加载数据
        self.load_data()
        
    def load_data(self):
        """根据模式加载图像或视频数据"""
        # 获取绝对路径
        if not os.path.isabs(self.data_path):
            self.data_path = os.path.join(os.path.dirname(__file__), self.data_path)
            
        if self.is_video_mode:
            self.load_video_data()
        else:
            self.load_image_data()
            
    def load_image_data(self):
        """加载图像文件列表"""
        # 支持的图像格式
        image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tiff', '*.tif', '*.webp']
        
        self.data_list = []
        for ext in image_extensions:
            pattern = os.path.join(self.data_path, ext)
            self.data_list.extend(glob.glob(pattern))
            # 同时搜索大写扩展名
            pattern_upper = os.path.join(self.data_path, ext.upper())
            self.data_list.extend(glob.glob(pattern_upper))
        
        # 排序确保一致的顺序
        self.data_list.sort()
        
        if not self.data_list:
            rospy.logerr(f"No image files found in {self.data_path}")
            rospy.loginfo("Supported formats: jpg, jpeg, png, bmp, tiff, tif, webp")
            rospy.signal_shutdown("No data found")
            return
            
        rospy.loginfo(f"Loaded {len(self.data_list)} image files")
        for i, img_path in enumerate(self.data_list[:5]):  # 显示前5个文件
            rospy.loginfo(f"  [{i}] {os.path.basename(img_path)}")
        if len(self.data_list) > 5:
            rospy.loginfo(f"  ... and {len(self.data_list) - 5} more")
            
    def load_video_data(self):
        """加载视频文件"""
        # 支持的视频格式
        video_extensions = ['*.mp4', '*.avi', '*.mov', '*.mkv', '*.wmv', '*.flv']
        
        video_files = []
        for ext in video_extensions:
            pattern = os.path.join(self.data_path, ext)
            video_files.extend(glob.glob(pattern))
            # 同时搜索大写扩展名
            pattern_upper = os.path.join(self.data_path, ext.upper())
            video_files.extend(glob.glob(pattern_upper))
            
        if not video_files:
            rospy.logerr(f"No video files found in {self.data_path}")
            rospy.loginfo("Supported formats: mp4, avi, mov, mkv, wmv, flv")
            rospy.signal_shutdown("No data found")
            return
            
        # 使用第一个找到的视频文件
        video_file = sorted(video_files)[0]
        rospy.loginfo(f"Loading video file: {os.path.basename(video_file)}")
        
        self.video_cap = cv2.VideoCapture(video_file)
        if not self.video_cap.isOpened():
            rospy.logerr(f"Failed to open video file: {video_file}")
            rospy.signal_shutdown("Failed to open video")
            return
            
        # 获取视频信息
        frame_count = int(self.video_cap.get(cv2.CAP_PROP_FRAME_COUNT))
        video_fps = self.video_cap.get(cv2.CAP_PROP_FPS)
        duration = frame_count / video_fps if video_fps > 0 else 0
        
        rospy.loginfo(f"Video info:")
        rospy.loginfo(f"  Frames: {frame_count}")
        rospy.loginfo(f"  Video FPS: {video_fps:.2f}")
        rospy.loginfo(f"  Duration: {duration:.2f}s")
        
    def publish_image_sequence(self):
        """发布图像序列"""
        if not self.data_list:
            return False
            
        # 加载当前图像
        image_path = self.data_list[self.current_index]
        cv_image = cv2.imread(image_path)
        
        if cv_image is None:
            rospy.logwarn(f"Failed to load image: {os.path.basename(image_path)}")
            self.current_index = (self.current_index + 1) % len(self.data_list)
            return True
            
        # 转换为ROS消息并发布
        try:
            image_msg = self.bridge.cv2_to_imgmsg(cv_image, "bgr8")
            image_msg.header.stamp = rospy.Time.now()
            image_msg.header.frame_id = "camera_frame"
            self.image_pub.publish(image_msg)
            
            rospy.loginfo(f"Published image [{self.current_index}/{len(self.data_list)}]: {os.path.basename(image_path)}")
            
        except Exception as e:
            rospy.logerr(f"Failed to publish image: {e}")
            
        # 更新索引
        self.current_index += 1
        if self.current_index >= len(self.data_list):
            if self.loop:
                self.current_index = 0
                # rospy.loginfo("Restarting image sequence")
            else:
                # rospy.loginfo("Image sequence completed")
                return False
                
        return True
        
    def publish_video_frames(self):
        """发布视频帧"""
        if self.video_cap is None:
            return False
            
        # 读取下一帧
        ret, frame = self.video_cap.read()
        
        if not ret:
            if self.loop:
                # 重新开始视频
                self.video_cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
                ret, frame = self.video_cap.read()
                rospy.loginfo("Restarting video")
                if not ret:
                    rospy.logerr("Failed to restart video")
                    return False
            else:
                rospy.loginfo("Video completed")
                return False
                
        # 转换为ROS消息并发布
        try:
            image_msg = self.bridge.cv2_to_imgmsg(frame, "bgr8")
            image_msg.header.stamp = rospy.Time.now()
            image_msg.header.frame_id = "camera_frame"
            self.image_pub.publish(image_msg)
            
            # 获取当前帧位置
            current_frame = int(self.video_cap.get(cv2.CAP_PROP_POS_FRAMES))
            total_frames = int(self.video_cap.get(cv2.CAP_PROP_FRAME_COUNT))
            rospy.loginfo(f"Published video frame [{current_frame}/{total_frames}]")
            
        except Exception as e:
            rospy.logerr(f"Failed to publish video frame: {e}")
            
        return True
        
    def run(self):
        """主运行循环"""
        rospy.loginfo("Virtual camera started")
        
        while not rospy.is_shutdown():
            try:
                if self.is_video_mode:
                    if not self.publish_video_frames():
                        break
                else:
                    if not self.publish_image_sequence():
                        break
                        
                self.rate.sleep()
                
            except rospy.ROSInterruptException:
                break
            except Exception as e:
                rospy.logerr(f"Error in main loop: {e}")
                break
                
        # 清理资源
        if self.video_cap is not None:
            self.video_cap.release()
            
        rospy.loginfo("Virtual camera stopped")


def main():
    try:
        camera = VirtualCamera()
        camera.run()
    except rospy.ROSInterruptException:
        rospy.loginfo("Virtual camera interrupted")
    except Exception as e:
        rospy.logerr(f"Virtual camera error: {e}")


if __name__ == '__main__':
    main()
